#!/usr/bin/env python
# -*- coding: utf8 -*-
import os
import sys


def get_project_root(dirname):
    filepath = os.path.abspath(__file__)
    if os.path.islink(filepath):
        filepath = os.readlink(filepath)
    index = filepath.rindex(dirname)
    root = filepath[:index + len(dirname)]
    return root

BASE_DIR = get_project_root('labor')
sys.path.insert(0, BASE_DIR)
os.chdir(BASE_DIR)

import six
import time
import socket
import xmlrpclib
from datetime import datetime
from optparse import OptionParser

from utils.loader import import_string
from utils.getconf import config_getter
from utils.shell import Executer as shell
from utils.format import JSON

options = None


def header():
    return datetime.now().strftime('%Y-%m-%d %H:%M:%S') + ' '


class RPCException(Exception):

    def __init__(self, host):
        self.message = "connect to supervisord {0}:9001 failed".format(host)

    def __str__(self):
        return self.message


class DrainException(Exception):
    pass


def parse_backends(backends):
    backend_classes = {}
    dealers = []
    for class_string in list(set([e[0] for e in backends])):
        cls = import_string(class_string)
        backend_classes[class_string] = cls

    for class_string, class_conf in backends:
        cls = backend_classes.get(class_string)
        dealer = cls(conf=class_conf)
        dealers.append(dealer)
        dealer.get_nodes()

    return dealers[0].rolling_nodes, dealers


class RollingService(object):

    def __init__(self, hosts, dealers, rolling_config):
        self.hosts = hosts
        self.dealers = dealers
        self.process = rolling_config.supervisor.group
        self.code_dir = rolling_config.code_dir

        self.local_hostname = socket.gethostname().split('.')[0]
        self.local_ip = socket.gethostbyname(self.local_hostname)

    def run(self):
        local_node = self.update_local()
        self.hosts.remove(local_node)
        compress_file = self.compress_code()

        for host in self.hosts:
            try:
                print header() + "-" * 30 + " rolling {0} ".format(host) + '-' * 30
                rpc_cli = self.get_supervisor(host)
                self.drain_node(host)
                print header() + "sleep 3s for node cool-down"
                time.sleep(3)
                self.copy_to_remote(host, compress_file)
                self.rolling(rpc_cli, host)
            except RPCException, e:
                print e.__str__(), "will do next"
                continue

        print header() + "clean files"
        os.remove(compress_file)

    def get_supervisor(self, host):
        server = xmlrpclib.Server('http://{0}:9001/RPC2'.format(host))
        try:
            server.supervisor.getState()
            return server
        except socket.error:
            raise RPCException(host)

    def update_local(self):
        if self.local_hostname in self.hosts:
            local_node = self.local_hostname
        elif self.local_ip in self.hosts:
            local_node = self.local_ip
        else:
            raise Exception("local server is not in backservers")

        print header() + "-" * 30 + " rolling {0} ".format(local_node) + '-' * 30
        rpc_cli = self.get_supervisor(local_node)
        self.drain_node(local_node)
        print header() + "sleep 3s for node cool-down"
        time.sleep(3)
        self.update_code()
        self.rolling(rpc_cli, local_node)
        return local_node

    def rolling(self, rpc_cli, node):
        try:
            self.stop_process(rpc_cli, self.process)
            self.start_process(rpc_cli, self.process)
            self.enable_node(node)
            time.sleep(1)
            return node
        except Exception, e:
            print e
            if self.is_process_running(rpc_cli, self.process):
                error = "update failed, re-enable node: {0}".format(node)
                self._error_output(error)
                self.enable_node(node)
                return node
            else:
                error = "{0} on {1} is not running, interupt rolling!".format(self.process, node)
                self._error_output(error)
                raise Exception("can't continue")

    def update_code(self):
        step = "update_code"
        print header() + '{:70}'.format(step),
        try:
            shell.call("git pull", cwd=self.code_dir)
            self._output_result(True)
        except shell.ExecException, e:
            self._output_result(False)
            raise e

        try:
            output = shell.call(
                """git log -n 1 --pretty --format='{"hash":"%H","subject":"%s","author":"%an","timestamp":%at}'""", cwd=self.code_dir)
            commit = output[0]
            jbody = JSON.loads(commit)
            print header() + u"commit subject: \033[36m{}\033[0m".format(jbody['subject'])
        except Exception, e:
            raise e

    def compress_code(self):
        step = "compress_code"
        print header() + '{:70}'.format(step),
        compress_file = '/tmp/{0}.tar.gz'.format(self.process)
        try:
            shell.call("tar zcf {0} .".format(compress_file), cwd=self.code_dir)
            self._output_result(True)
            print header() + "compress to " + compress_file
            return compress_file
        except shell.ExecException, e:
            self._output_result(False)
            raise e

    def copy_to_remote(self, dest_host, filename):
        step = "copy_to_remote"
        print header() + '{:70}'.format(step),
        try:
            shell.call("""cat {0} | ssh {1} 'tar zxf - -C {2}'""".format(filename, dest_host, self.code_dir))
            self._output_result(True)
        except shell.ExecException:
            self._output_result(False)

    def drain_node(self, host):
        if options.do_drain:
            for dealer in self.dealers:
                step = "drain_node on {0}: {1}".format(dealer.label, host)
                print header() + '{:70}'.format(step),
                success = dealer.drainNode(host)
                self._output_result(success)
                if not success:
                    raise DrainException("drain_node failed")

    def enable_node(self, host):
        if options.do_drain:
            for dealer in self.dealers:
                is_vaild = self.vaild_node(dealer, host)
                step = "enable_node on {0}: {1}".format(dealer.label, host)
                print header() + '{:70}'.format(step),
                if is_vaild:
                    success = dealer.enableNode(host)
                    self._output_result(success)
                else:
                    self._output_result(False)

    def vaild_node(self, dealer, host):
        retry = 5
        step = "check_health on {0}: {1}".format(dealer.label, host)
        for i in xrange(retry):
            print header() + '{:70}'.format(step),
            health = dealer.checkHealth(host)
            self._output_result(health)
            if not health:
                time.sleep(3)
            else:
                return True

        return False

    def stop_process(self, rpc_cli, process):
        if options.do_restart:
            step = "stop_process: {0}".format(process)
            print header() + '{:70}'.format(step),
            success = rpc_cli.supervisor.stopProcess(process)
            self._output_result(success)

    def start_process(self, rpc_cli, process):
        if options.do_restart:
            step = "start_process: {0}".format(process)
            print header() + '{:70}'.format(step),
            success = rpc_cli.supervisor.startProcess(process)
            self._output_result(success)

    def is_process_running(self, rpc_cli, process):
        process_info = rpc_cli.supervisor.getProcessInfo(process)
        return process_info["pid"]

    def _output_result(self, success):
        if success:
            print "\033[32m[SUCCESS]\033[0m"
        else:
            print "\033[31m[FAILED]\033[0m"

    def _error_output(self, msg):
        print "\033[31m{0}\033[0m".format(msg)


def change_user(user=None):
    from pwd import getpwnam
    if user is not None:
        uid, gid = getpwnam(user)[2:4]
    else:
        uid = os.getuid()
        gid = os.getgid()
    os.setgid(gid)
    os.setuid(uid)


def getopts():
    usage = "usage: %prog -p <process>"
    parser = OptionParser(usage=usage)
    parser.add_option("-p", "--process", action="store", type="string", dest="process", help="rolling the process")
    parser.add_option("-l", "--list", action="store_true", dest="show_list", default=False, help="list arrival processes")
    parser.add_option("--nodrain", action="store_false", dest="do_drain", default=True)
    parser.add_option("--norestart", action="store_false", dest="do_restart", default=True)
    global options
    (options, args) = parser.parse_args()
    if options.show_list:
        process_list = config_getter.get('rolling').keys()
        for p in process_list:
            print p
        os._exit(0)
    elif options.process is None:
        parser.print_help()
        os._exit(1)


def main():
    getopts()

    process = options.process
    rolling_config = config_getter.get('rolling.{0}'.format(process))
    backends = rolling_config.backends
    hosts, dealers = parse_backends(backends)
    print header() + "rolling nodes for process '{0}': {1}".format(options.process, ', '.join(hosts))
    roller = RollingService(hosts, dealers, rolling_config)
    roller.run()

if __name__ == '__main__':
    change_user('rain')
    main()
